from torchvision.datasets import FashionMNIST
from torchvision import transforms
import torch.utils.data as Data
import numpy as np
import matplotlib.pyplot as plt

train_data = FashionMNIST('./data',
                          train=True,
                          transform=transforms.Compose([transforms.Resize(size=224),transforms.ToTensor()]),
                          download=True)

train_loader = Data.DataLoader(dataset=train_data,
                               batch_size=64,
                               shuffle=True,#shuffle是打乱图片
                               num_workers=0)

#获得一个Batch的数据,
# 结果：['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']
for step,(b_x,b_y) in enumerate(train_loader):
#train_loader 返回一个批次的数据，其中 b_x 是输入数据（通常是图像），b_y 是对应的标签。
    if step>0:#索引 step 是否大于 0。如果是，则跳出循环。这意味着循环只会执行一次，处理第一个批次的数据。
        break
#一个批次图片是64*28*28*1,则通过squeeze后是28*28*1,移除第一维度
batch_x=b_x.squeeze().numpy()#squeeze 方法用于去除 b_x 中所有大小为 1 的维度。这通常用于简化张量的形状。转为numpy数组方便画图
batch_y=b_y.numpy()
class_label=train_data.classes
print(class_label)
print(batch_y)#[2 4 4 5 2 6 4 0 6 3 4 7 3 3 6 3 0...
#print(batch_x[0])  # 查看第一个样本的内容
print("The size of batch_x is: ",batch_x.shape)


#可视化一个Batch的图集
plt.figure(figsize=(12,5))
for ii in np.arange(len(batch_y)):
    plt.subplot(4,16,ii+1)#plt.subplot(nrows, ncols, index, **kwargs)
    plt.imshow(batch_x[ii,:,:],cmap=plt.cm.gray)
    plt.title(class_label[batch_y[ii]],size=10)
    plt.axis('off')
    plt.subplots_adjust(wspace=0.05)
plt.show()

